package main

import "fmt"

func main() {
	fmt.Println(wordBreak("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaab", []string{"aa", "a"}))
	fmt.Println(wordBreak("leetcode", []string{"leet", "code"}))
}

func wordBreak(s string, wordDict []string) bool {
	trie := constructTrie(wordDict)
	//dfs
	dp := make([]bool, len(s)+1)
	dp[0] = true
	for i := 1; i < len(dp); i++ {
		for j := i - 1; j >= 0; j-- {
			if trie.search(s[j:i]) && dp[j] {
				dp[i] = true
				continue
			}
		}
	}
	return dp[len(dp)-1]
}

type trieNode struct {
	next    map[rune]*trieNode
	isLeave bool
}

func newTrieNode() *trieNode {
	return &trieNode{
		next: make(map[rune]*trieNode),
	}
}

func constructTrie(wordDict []string) *trieNode {
	root := newTrieNode()
	for _, v := range wordDict {
		root.insert(v)
	}
	return root
}

func (trie *trieNode) insert(val string) {
	tmpTrie := trie
	for _, v := range val {
		if tmpTrie.next[v] == nil {
			tmpTrie.next[v] = newTrieNode()
		}
		tmpTrie = tmpTrie.next[v]
	}
	tmpTrie.isLeave = true
}

func (trie *trieNode) search(val string) bool {
	tmpTrie := trie
	for _, v := range val {
		if tmpTrie.next[v] == nil {
			return false
		}
		tmpTrie = tmpTrie.next[v]
	}
	return tmpTrie.isLeave
}
